import os
import cv2
import numpy as np
import pandas as pd
import glob
from pathlib import Path
import json
import sys
sys.path.append('..')
from utils.json import json_to_image, mask_to_json

# def my_train_test_split(data, test_size=0.1, by='code', random_state=752):
#     codes = np.unique(data[by])
#     train_df = None
#     val_df = None
#     for code in codes:
#         tmp = data[data[by]==code].reset_index(drop=True)
#         tmp = tmp.sample(frac=1, random_state=random_state).reset_index(drop=True)
#         len_train = int(len(tmp) * (1-test_size))
#         if len_train == 0:
#             train_df = pd.concat([train_df, tmp[:]])
#         else:
#             train_df = pd.concat([train_df, tmp[:len_train]])
#             val_df = pd.concat([val_df, tmp[len_train:]])
#     return train_df, val_df


# def search_json(img_path):
#     json_path = Path(img_path).with_suffix('.json')
#     if json_path.exists():
#         return str(json_path)
#     return '0'

# def code_limit_exec(df, memo):
#     for code in memo:
#         df1 = df[df['code'] != code]
#         df2 = df[df['code'] == code]
#         if len(df2) > memo[code]:
#             df2 = df2.sample(frac=1, random_state=10).reset_index(drop=True).iloc[:memo[code]]
#             df = pd.concat([df1, df2]).reset_index(drop=True)
#     return df


# ########################################
# path = r'/home/hanyu/TCL/AutoRepair_AA_T7/t7_training_data/all_img_cut_pad'
# code_num_limits = {'TGXID':800}

# df = pd.DataFrame()

# img_paths = glob.glob(os.path.join(path, '*/*.jpg'))
# # label_paths = glob.glob(os.path.join(path, '*/*.json'))

# df['image'] = img_paths
# df['code'] = df['image'].apply(lambda x: str(Path(x).parent.name))
# df = code_limit_exec(df, code_num_limits)
# df['label'] = df['image'].apply(lambda x: search_json(x))

# train_df, val_df = my_train_test_split(df, 0.15)

# os.makedirs('./datasets/csv', exist_ok=True)
# train_df.to_csv(r'./datasets/csv/AA_T7_train.csv', index=False)
# val_df.to_csv(r'./datasets/csv/AA_T7_val.csv', index=False)



# def json_filter(json_path, dst_path):
#     with open(json_path, 'r') as f:
#         content = f.read()
#         data = json.loads(content)
    
#     data['imageData'] = 'null'
    
#     new_data_shapes= []
#     for d in data['shapes']:
#         d['label'] = d['label'].replace(' ', '')
#         if len(d['points']) < 3:
#             continue
#         if len(d['label']) == 6 and d['label'][-1] == '1':  # 去掉前程修补
#             continue
#         if len(d['label']) == 6:  # 合并2 & 3至主code
#             d['label'] = d['label'][:-1]
#         new_data_shapes.append(d)
#     data['shapes'] = new_data_shapes

#     json_name = str(Path(json_path).name)
#     code = str(Path(json_path).parent.name)
#     json_dst_path = os.path.join(dst_path, code, json_name)
#     os.makedirs(Path(json_dst_path).parent, exist_ok=True)
#     with open(json_dst_path, 'w') as f:
#         b = json.dumps(data, indent=4)
#         f.write(b)

def get_img_path(json_path, code, img_dst_root):
    img_name = str(Path(json_path).name)[:-4] + 'jpg'
    return os.path.join(img_dst_root, code, img_name)

def one_img_cut(img, mask):
    h, w, c = img.shape
    mid_h = h//2
    mid_w = w//2

    new_imgs, new_masks = [], []
    for dh, dw in [[0, 0], [0, 1], [1, 1], [1, 0]]:
        new_imgs.append(img[dh*mid_h:(dh+1)*mid_h, dw*mid_w:(dw+1)*mid_w])
        new_masks.append(mask[dh*mid_h:(dh+1)*mid_h, dw*mid_w:(dw+1)*mid_w])
    
    return new_imgs, new_masks


# # json_to_image, mask_to_json
def img_cut(json_path, img_path, code, no_defect_code, dst_path):
    img_stem = str(Path(img_path).stem)
    
    img = cv2.imread(img_path, 1)
    mask = json_to_image(img.shape, json_path)

    cutted_imgs, cutted_masks = one_img_cut(img, mask)
    for i, (c_img, c_mask) in enumerate(zip(cutted_imgs, cutted_masks)):
        # print('c mask shape: ', c_mask.shape)
        if 1 not in c_mask:
            tmp_code = 'X' + no_defect_code[1:]
        else:
            tmp_code = 'X' + code[1:] 

        img_name = img_stem + '_%d.jpg'%(i+1)
        mask_name = img_stem + '_%d.json'%(i+1)
        img_dst_path = os.path.join(dst_path, 'images', tmp_code)
        mask_dst_path = os.path.join(dst_path, 'annotations', tmp_code)

        os.makedirs(img_dst_path, exist_ok=True)
        cv2.imwrite(os.path.join(img_dst_path, img_name), c_img)

        if 1 in c_mask:
            os.makedirs(mask_dst_path, exist_ok=True)
            mask_json = mask_to_json(mask=c_mask, imagePath=img_name, label_name=tmp_code)
            with open(os.path.join(mask_dst_path, mask_name), 'w') as f:
                b = json.dumps(mask_json, indent=4)
                f.write(b)


# TDSBN TTSPT TCNFO  TTPRP TCSBP
def main():
    # parameters
    ################################################################ 
    img_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_img_cut_pad'
    json_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/filtered_json'
    dst_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/cutted_data'

    need_cut_codes = ['TDSBN', 'TTSPT', 'TCNFO', 'TTPRP', 'TCSBP']
    no_defect_code = 'TGXID'
    ################################################################

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    df = pd.DataFrame({'json':glob.glob(os.path.join(json_path, '*/*.json'))})
    df['code'] = df['json'].apply(lambda x:str(Path(x).parent.name))
    df = df[df['code'].isin(need_cut_codes)].reset_index(drop=True)
    if is_pandarallel:
        df['img_path'] = df.parallel_apply(lambda x: get_img_path(x['json'], x['code'], img_path), axis=1)
    else:
        df['img_path'] = df.apply(lambda x: get_img_path(x['json'], x['code'], img_path), axis=1)
    
    print('len df: ', len(df))
    
    if is_pandarallel:
        df.parallel_apply(lambda x: img_cut(json_path=x['json'], img_path=x['img_path'], code=x['code'], no_defect_code=no_defect_code, dst_path=dst_path), axis=1)
    else:
        df.apply(lambda x: img_cut(json_path=x['json'], img_path=x['img_path'], code=x['code'], no_defect_code=no_defect_code, dst_path=dst_path), axis=1)


if __name__=='__main__':
    main()
